from pydantic import BaseModel, Field
from typing import Optional, Any
from enum import Enum


class TaskStatus(str, Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    NOT_FOUND = "not_found"


class OcrLlmModel(BaseModel):
    """OCR+LLM任务的数据模型"""

    task_id: Optional[str] = Field(default="", description="任务ID，唯一标识符")
    scenario: str = Field(description="场景代码")

    def replace(self, **changes):
        """返回一个新的 OcrTaskResult 实例，替换指定的字段值"""
        updated_data = self.model_dump()
        updated_data.update(changes)
        return self.__class__(**updated_data)


class OcrTask(OcrLlmModel):
    """
    OcrTask 数据模型，表示一个 OCR（光学字符识别）任务。
    属性:
        task_id（str）: 任务ID
        scenario (str): 场景代码
        image_path (str): 待识别图片的路径
        llm (bool): 是否需要大模型处理，默认为 False
        user_prompt: 用户提示词，JSON格式
        extra：额外的参数支持，JSON格式
    """

    """OCR任务的数据模型"""
    image_path: str = Field(description="待识别图片的路径")
    optimized_path: str = Field(description="优化后图片路径")
    llm: bool = Field(default=False, description="是否需要大模型处理")
    user_prompt: str = Field(description="用户提示词，JSON格式")
    extra: str = Field(description="额外的参数支持，JSON格式")


class OcrTaskResult(OcrLlmModel):
    """
    OcrTaskResult 数据模型用于表示 OCR 任务的结果。
    属性:
        task_id (Optional[str]): 任务唯一标识符，默认为空字符串。
        scenario (str): 场景代码
        texts (list[str]): 识别出的文本列表
        img_ms (Optional[int]): 图像预处理耗时（毫秒），默认为 0
        ocr_ms (Optional[int]): OCR 识别耗时（毫秒），默认为 0
        status (TaskStatus): 任务状态，默认为 PENDING
        error (Optional[str]): 错误信息（如果有），默认为空字符串
    """

    """任务结果的数据模型"""
    texts: list[str] = Field(description="识别出的文本列表", default=[])
    img_ms: Optional[int] = Field(default=0, description="优化处理耗时，单位毫秒")
    ocr_ms: Optional[int] = Field(default=0, description="OCR处理耗时，单位毫秒")
    status: TaskStatus = Field(
        default=TaskStatus.PENDING, description="任务状态，默认为 pending"
    )
    error: Optional[str] = Field(default="", description="错误信息，如果有的话")


class LlmTask(OcrLlmModel):
    """
    LlmTask 数据模型用于表示一个与大语言模型（LLM）相关的任务。
    属性:
        task_id（str）: 任务ID
        scenario (str): 场景代码
        ocr_text (list[str]): 识别出的文本列表
        user_prompt (str): 用户输入的提示信息，用于指导LLM生成响应
    """

    """任务数据模型"""
    ocr_text: str = Field(description="OCR识别的文字")
    user_prompt: str = Field(description="用户提示词")


class LlmTaskResult(OcrLlmModel):
    """
    LlmTaskResult 数据模型用于表示一次 LLM 任务的结果。
    属性:
        task_id (Optional[str]): 任务唯一标识符，默认为空字符串。
        scenario (str): 场景代码
        data (list[str]): 任务处理得到的结果数据列表，默认为空列表。
        llm_ms (Optional[int]): LLM 处理耗时，单位为毫秒，默认为 0。
        status (TaskStatus): 任务状态，默认为 PENDING。
        error (Optional[str]): 错误信息（如果有），默认为空字符串。
    """

    """任务数据模型"""

    data: Any = Field(description="LLM处理结果数据列表", default=[])
    llm_ms: Optional[int] = Field(default=0, description="LLM处理耗时，单位毫秒")
    status: TaskStatus = Field(
        default=TaskStatus.PENDING, description="任务状态，默认为 pending"
    )
    error: Optional[str] = Field(default="", description="错误信息，如果有的话")


class OcrLlmTaskResult(OcrLlmModel):
    """
    LlmTaskResult 数据模型用于表示一次 LLM 任务的结果。
    属性:
        task_id (Optional[str]): 任务唯一标识符，默认为空字符串。
        scenario (str): 场景代码
        texts (list[str]): 识别出的文本列表
        data (list[str]): 任务处理得到的结果数据列表，默认为空列表
        img_ms (Optional[int]): 图像预处理耗时（毫秒），默认为 0
        llm_ms (Optional[int]): LLM 处理耗时，单位为毫秒，默认为 0
        status (TaskStatus): 任务状态，默认为 PENDING。
        error (Optional[str]): 错误信息（如果有），默认为空字符串。
    """

    """任务数据模型"""
    texts: list[str] = Field(description="识别出的文本列表", default=[])
    data: list = Field(description="LLM处理结果数据列表", default=[])
    img_ms: Optional[int] = Field(default=0, description="图片优化处理耗时，单位毫秒")
    llm_ms: Optional[int] = Field(default=0, description="LLM处理耗时，单位毫秒")
    status: TaskStatus = Field(
        default=TaskStatus.PENDING, description="任务状态，默认为 pending"
    )
    error: Optional[str] = Field(default="", description="错误信息，如果有的话")
